
__version__ = '0.0.1'
__authors__ = ['Ozgur Yilmaz']


import pymel.core as pm
import oyAuxiliaryFunctions as oyAux

class LimbRigger(object):
    """the base class that will help to create the rig
    and hold the data about the process
    """
    
    def __init__(self, **kwords ):
        """oyAutoRig \n
        useHires : to setup a hires joint setup
        useSplit : adds an extra joint in the middle of the joint hierarchy
        useRefineControls : adds refine joints, implies useHires
        setupForFoot : setups for foot
        isStretchy : setups a stretchy limb
        upperSplitCnt : sets the joint count in the upper part of the limb
        lowerSplitCnt : set the joint count in the lower part of the limb
        """
        
        self.__version__ = '0.0.1'

        # parse keyword arguments
        self.parseArgs(**kwords)

        self._refineJoints_radius = 0.25
        self._hiresJoints_radius = 0.5
        self._controllerJoints_radius = 1.0
        self._ikJoints_radius = 1.5
        self._fkJoints_radius = 2.0
        
        
        # nodes
        #self._baseJoint = pm.Joint
        #self._ikFKSwitchController = pm.Joint
        #self._ikPoleController = pm.Joint
        #self._ikEndController = pm.Joint
        
        #self._refineJoints = [pm.Joint] * 0
        #self._hiresJoints = [pm.Joint] * 0
        #self._controllerJoints = [pm.Joint] * 0
        #self._ikJoints = [pm.Joint] * 0
        #self._fkJoints = [pm.Joint] * 0

        self._baseJoint = ''
        self._ikFKSwitchController = ''
        self._ikPoleController = ''
        self._ikEndController = ''
        
        self._refineJoints = [pm.Joint] * 0
        self._hiresJoints = [pm.Joint] * 0
        self._controllerJoints = [pm.Joint] * 0
        self._ikJoints = [pm.Joint] * 0
        self._fkJoints = [pm.Joint] * 0
        
    
    def checkArgument(self, longName, shortName, defaultValue, **kwords):
        if longName in kwords:
            returnArg = kwords[longName]
        elif shortName in kwords:
            returnArg = kwords[shortName]
        else:
            returnArg = defaultValue
        return returnArg
    

    def parseArgs(self, **kwords):
        # parse keywords
        
        #verbose = True
        
        # ------------------------------
        # BASE PARAMETERS
        # ------------------------------
        
        # useHires or uh
        self._useHires = self.checkArgument('useHires', 'uh', True, **kwords)
        
        # useSplit or us
        self._useSplit = self.checkArgument('useSplit', 'us', True, **kwords)
        
        # setupForFoot or sff
        self._setupForFoot = self.checkArgument('setupForFoot', 'sff', False, **kwords)
        
        # isStretchy or iss
        self._isStretchy = self.checkArgument('isStretchy', 'iss', True, **kwords)
        
        # upperSplitCnt or usc
        self._upperSplitCnt = self.checkArgument('upperSplitCnt', 'usc', 2, **kwords)
        
        # lowerSplitCnt or lsc
        self._lowerSplitCnt = self.checkArgument('lowerSplitCnt', 'lsc', 2, **kwords)
        
        # ------------------------------
        # JOINT NAMING
        # ------------------------------
        
        # baseName or bn
        self._baseName = self.checkArgument('baseName', 'bn', 'baseName', **kwords)
        
        # refineLabel or rl
        self._refineJoints_label = self.checkArgument('refineLabel', 'rl', 'refineJoint', **kwords)
        
        # hiresLabel or hl
        self._hiresJoints_label = self.checkArgument('hiresLabel', 'hl', 'hiresJoint', **kwords)
        
        # controllerLabel or cl
        self._controllerJoints_label = self.checkArgument('controllerLabel','cl', 'controllerJoint', **kwords)
        
        # IKLabel or ikl
        self._ikJoints_label = self.checkArgument('ikLabel', 'ikl', 'IKJoint', **kwords)
        
        # FKLabel or fkl
        self._fkJoints_label = self.checkArgument('fkLabel', 'fkl', 'FKJoint', **kwords)
        
        # sideLabel or sl
        self._sideLabel = self.checkArgument('sideLabel', 'sl', '', **kwords)

        # useRadius or ur
        self._useRadius = self.checkArgument('useRadius', 'ur', True, **kwords)
        
        # verbose or v
        self._verbose = self.checkArgument('verbose', 'v', False, **kwords)
        
        if self._verbose:
            space = 20
            print '*' * (space * 2)
            print 'oyRigTools.oyLimbRigger'.center(space*2)
            print '-' * (space * 2)
            print 'useSplit'.ljust(space),': %s' % self._useSplit
            print 'useHires'.ljust(space),': %s' % self._useHires
            print 'isStretchy'.ljust(space),': %s' % self._isStretchy
            print 'setupForFoot'.ljust(space),': %s' % self._setupForFoot
            print 'lowerSplitCnt'.ljust(space),': %s' % self._lowerSplitCnt
            print 'upperSplitcnt'.ljust(space),': %s' % self._upperSplitCnt
            print 'baseName'.ljust(space),': %s' % self._baseName
            print 'hiresLabel'.ljust(space),': %s' % self._hiresJoints_label
            print 'controllerLabel'.ljust(space),': %s' % self._controllerJoints_label
            print 'ikLabel'.ljust(space),': %s' % self._ikJoints_label
            print 'fkLabel'.ljust(space),': %s' % self._fkJoints_label
            print 'sideLabel'.ljust(space),': %s' % self._sideLabel
            print 'useRaidus'.ljust(space),': %s' % self._useRadius
            print '*' * (space * 2)
    
    
    def set_baseJoint(self, baseJoint):
        """sets the base joint
        """
        if isinstance(baseJoint, str):
            try:
                self._baseJoint = pm.Joint( baseJoint )
            except pm.MayaNodeError:
                print 'error: no such joint : %s' % baseJoint
        
        elif isinstance(baseJoint, pm.Joint):
            self._baseJoint = baseJoint
        
    def set_ikFkSwitchController(self, ikFkSwitchController):
        """sets the ikFkSwitchController
        """
        if isinstance(ikFkSwitchController, str):
            try:
                self._ikFKSwitchController = pm.Transform( ikFkSwitchController )
            except pm.MayaNodeError:
                print 'error: no such node : %s' % ikFkSwitchController
                
        elif isinstance(ikFkSwitchController, pm.Transform):
            self._ikFKSwitchController = ikFkSwitchController
    
    def set_ikPoleController(self, ikPoleController ):
        """sets the ikPoleController
        """
        if isinstance(ikPoleController, str):
            try:
                self._ikPoleController = pm.Transform( ikPoleController )
            except pm.MayaNodeError:
                print 'error: no such node : %s' % ikPoleController
                
        elif isinstance(ikPoleController, pm.Transform):
            self._ikPoleController = ikPoleController
    
    def set_ikEndController(self, ikEndController ):
        """sets the ikEndController
        """
        if isinstance(ikEndController, str):
            try:
                self._ikEndController = pm.Tranform( ikEndController )
            except pm.MayaNodeError:
                print 'error: no such node : %s' % ikEndController
                
        elif isinstance(ikEndController, pm.Transform):
            self._ikEndController = ikEndController

    def insert_parent_joints(self, joint, jointCnt):
        """inserts the amount of joints between the joint and the parent of that
        joint
        
        returns a list that contains the only the newly created joints
        """
        
        joint = self.getValidDagNode( joint )
        
        # get the local transforms
        locTra = joint.getAttr( 't' )
        
        # get the parent
        parent = self.getParent( joint )
        
        print 'parent count : %s' % len(parent)
        if len(parent) is 0:
            return
        
        newJoints = [pm.DagNode] * (jointCnt + 1)
        newJoints[0] = parent[0]
        
        # duplicate the parent jointCnt times
        # and delete all the childs of the new joints
        for i in range(1, jointCnt + 1):
            # duplicate the parent
            newJoints[i] = pm.duplicate( parent )[0]
            
            # get the children and delete them
            children = self.getChildren( newJoints[i] )
            pm.delete( children )
            
            # parent the newJoints[i] under the previous
            pm.parent ( newJoints[i], newJoints[i-1])
            
            # move the new parent to the new position
            newJoints[i].setAttr( 't', locTra / float( jointCnt + 1 ) )
            
            # connect the parents scale to the inverseScale of the new joint
            newJoints[i-1].scale >> newJoints[i].inverseScale
        
        # finally parent the current joint to the last joint
        pm.parent( joint, newJoints[-1] )
        
        # connect the scale to inverseScale
        newJoints[-1].scale >> joint.inverseScale
        
        # remove the parent joint from the list
        newJoints.pop(0)
        
        return newJoints
    
    def insert_joints(self, joint, jointCnt ):
        '''works like insert_parent_joint but inserts towards the child
        '''
        
        allNewJoints = [pm.Joint] * 0
        
        # get all the children
        children = pm.listRelatives( joint, c=True )
        
        # run insert_joint for every child
        for child in children:
            if type(child) is pm.Joint:
                newJoints = self.insert_parent_joints( child, jointCnt )
                self.append_to_list( newJoints, allNewJoints )
        
        return allNewJoints
        
    def append_to_list(self, list1, list2 ):
        '''appends all the elements from the list1 to list2
        '''
        list3 = list2
        for element in list1:
            list3.append( element )
        
        return list3
        
    def getChildren(self, node):
        dagNode = self.getValidDagNode( node )
        if dagNode is not None:
            return pm.listRelatives ( dagNode, c=True )
        
    
    def getParent(self, node):
        #dagNode = self.getValidDagNode( node )
        #if isinstance(dagNode, pm.DagNode):
            #return pm.listRelatives ( dagNode, p=True )
        return pm.listRelatives( node , p=True )
        
        
    def getValidDagNode(self, node):
        nodeName = ''
        try:
            dagNode = pm.DagNode( node )
        except pm.MayaNodeError:
            if isinstance(node,str):
                nodeName = node
            elif isinstance(node, pm.DagNode):
                nodeName = node.name()
            
            print 'Error: no node named : %s' % nodeName
            return None
        
        return dagNode
    

    def create_hiresJoints(self):
        '''creates the hires joint hierarchy from the baseJoint group
        '''
        
        # duplicate the baseJoint hirerarchy
        hiresJoints = self.duplicate_baseJoint()
        
        # split the hires joints [0] and [1] or [2] if useSplit is set
        # split the upper joints
        newUpperJoints = self.insert_joints( hiresJoints[0] , self._upperSplitCnt )

        if self._useSplit:
            splitStartIndex = 2
        else:
            splitStartIndex = 1
        
        # split the lower joints
        newLowerJoints = self.insert_joints( hiresJoints[ splitStartIndex ] , self._lowerSplitCnt )
        
        # insert the newLowerJoints to hiresJoints list
        for i in range(len(newLowerJoints)-1,-1,-1):
            hiresJoints.insert( splitStartIndex, newLowerJoints[i] )
        
        # insert the newUppoerJoints to hiresJoints list
        for i in range(len(newUpperJoints)-1,-1,-1):
            hiresJoints.insert( 1, newUpperJoints[i] )
        
        # resize the joints
        if self._useRadius:
            self.resize_joints( hiresJoints, self._hiresJoints_radius )
        
        # rename the joints
        self.rename_joints( hiresJoints, self._hiresJoints_label)
        
        return hiresJoints
    
    def create_controllerJoints(self):
        '''creates the controllerJoint hierarchy
        '''
        # just duplicate the baseJoints
        controllerJoints = self.duplicate_baseJoint()
        
        # resize the joints
        self.resize_joints( controllerJoints, self._controllerJoints_radius )
        
        # rename the joints
        self.rename_joints( controllerJoints, self._controllerJoints_label )
        
        return controllerJoints

    
    def create_ikJoints(self):
        '''creates the ikJoint hierarchy
        '''
        # just duplicate the baseJoints
        ikJoints = self.duplicate_baseJoint()
        
        # resize the joints
        self.resize_joints( ikJoints, self._ikJoints_radius )
        
        # rename the joints
        self.rename_joints( ikJoints, self._ikJoints_label )
        
        return ikJoints
    
    def create_fkJoints(self):
        '''creates the fkJoint hierarchy
        '''
        # just duplicate the baseJoints
        fkJoints = self.duplicate_baseJoint()
        
        # resize the joints
        self.resize_joints( fkJoints, self._fkJoints_radius )
        
        # rename the joints
        self.rename_joints( fkJoints, self._fkJoints_label )
        
        return fkJoints

    def resize_joints(self, joints, size):
        for joint in joints:
            joint.setAttr('radius', size )
    
    def rename_joints(self, joints, jointGrpLabel):
        '''resizes joints by using the radius attribute
        '''
        # baseLabel + '_' + sideLabel + '_' + hiresLabel + '_' + counter
        for joint in joints:
            newName = self._baseName + '_' + self._sideLabel + '_' + jointGrpLabel + '#'  
            pm.rename( joint, newName )
    
    def duplicate_baseJoint(self):
        '''creates a duplicate of the neccessary part of the baseJoint
        '''
        
        # calculate the jointCnt
        if self._useSplit:
            jointCnt = 4
        else:
            jointCnt = 3
        
        # start from the baseJoint itself
        joint = self._baseJoint
        
        newJoint = [pm.Joint] * 0
        
        for i in range(0,jointCnt):
            # duplicate the joint
            dupJoint = pm.duplicate( joint )[0]
            
            # delete all the children of the new joint
            pm.delete ( pm.listRelatives( dupJoint, c=True ) )
            
            # append it to the new joints list
            newJoint.append( dupJoint )
                        
            # move forward in the hierarchy
            if i < jointCnt-1:
                children = pm.listRelatives( joint, c=True  )
                joint = children[0]
        
        # parent the newJoints to each other
        for i in range(jointCnt-1,0,-1):
            pm.parent( newJoint[i], newJoint[i-1] )
        
        # unparent the first joint in the list
        if len(pm.listRelatives( newJoint[0], p=True)) > 0:
            pm.parent( newJoint[0], w=True )
        
        return newJoint
    
    
    
    def setupLimb(self, *args, **kwords ):
        # create the joint groups
        self._hiresJoints = self.create_hiresJoints()
        self._controllerJoints = self.create_controllerJoints()
        self._ikJoints = self.create_ikJoints()
        self._fkJoints = self.create_fkJoints()

        #ctrlJntCnt = len( self._controllerJoints )
        
        ## orient constrain ik and fk joints to controller joint
        #for i in range(0,ctrlJntCnt):
            #orientCnts = pm.orientConstraint ( self._ikJoints[i], self._controllerJoints[i] )
            #pm.orientConstraint ( self._fkJoints[i], self._controllerJoints[i] )
            
            #self.setConstraintInterpolation( orientCnts )
        
    def distributeRotation(self, sourceObjects, targetObjects):
        '''creates an aimConstraint setup between sourceObjects to targetObjects
        to distribute the rotations smoothly
        
        source objects are the objects at the begining and the end
        
        target objects are the objects between them
        '''
        
        w1 = len(targetObjects) - 1
        w2 = 0

        # create upLocator for sourceObjects
        sBaseLoc1, sUpLoc1 = self.create_upLocators( sourceObjects[0] )
        sBaseLoc2, sUpLoc2 = self.create_upLocators( sourceObjects[1] )
        
        for tObj in targetObjects:
            tObj = self.getValidDagNode(tObj)
            
            # create the up locators
            tBaseLoc, tUpLoc = self.create_upLocators( tObj )
            
            tObjRotOr = rotateOrder( tObj.getAttr('ro') )
            
            aimVect = tObjRotOr.firstAxisList
            upVect = tObjRotOr.secondAxisList
            
            # create an aim constraint towards the second sourceObject
            aimC = pm.aimConstraint( sourceObjects[1], tObj, aim=aimVect, mo=False, wut='object', wuo=tUpLoc )
            
            # creaet a point constraint between the source objects upLoc and
            # the current target objects upLocator
            pC = pm.pointConstraint( sUpLoc1, tUpLoc, w=w1)
            pm.pointConstraint( sUpLoc2, tUpLoc, w=w2)
            
            w1 -= 1
            w2 += 1
            
    
    
    def create_upLocators(self, obj):
        '''creates a locator under the parent of the obj
        '''
        
        obj = self.getValidDagNode(obj)
        
        # create the locators
        baseLocator = pm.spaceLocator()
        pm.rename( baseLocator, obj.name() + '_upVector_baseLoc#' )
        
        upLocator = pm.spaceLocator()
        pm.rename( upLocator, obj.name() + '_upVector#')
        
        # parent the upLocator under baseLocator
        pm.parent( upLocator, baseLocator )
        
        # zero-out transformations of upLocator
        pm.parent( baseLocator, obj, r=True)
        
        # parent baseLocator under the parent of the obj
        parent = pm.listRelatives(obj, p=True)
        if len(parent) > 0:
            pm.parent( baseLocator, parent )
        else:
            pm.parent( baseLocator, world=True )

        # get the rotate order of the obj
        objRotOr = rotateOrder( obj.getAttr('ro') )
       
        # get the second-axis
        secondAxis = objRotOr.secondAxis
        
        # get the axis array
        axisArray = objRotOr.secondAxisList
        
        # move the upLocator with axisArray
        upLocator.setAttr('t',axisArray)
        
        return baseLocator, upLocator
    
    
    
    def setConstraintInterpolation(self, constraint ):
        constraint.setAttr('interpType',2)

class rotateOrder(object):
    def __init__(self, orderIndex):
        self.order = dict()
        self.axis = dict()
        
        self.order['XYZ'] = '0'
        self.order['YZX'] = '1'
        self.order['ZXY'] = '2'
        self.order['XZY'] = '3'
        self.order['YXZ'] = '4'
        self.order['ZYX'] = '5'
        
        self.order['0'] = 'XYZ'
        self.order['1'] = 'YZX'
        self.order['2'] = 'ZXY'
        self.order['3'] = 'XZY'
        self.order['4'] = 'YXZ'
        self.order['5'] = 'ZYX'
        
        self.axis['X'] = [1,0,0]
        self.axis['Y'] = [0,1,0]
        self.axis['Z'] = [0,0,1]
        
        self.axis['-X'] = [-1,0,0]
        self.axis['-Y'] = [0,-1,0]
        self.axis['-Z'] = [0,0,-1]
        
        self.currentOrder = self.order[ str(int(orderIndex)%3) ]
        
        self.firstAxis = self._getAxis(0)
        self.secondAxis = self._getAxis(1)
        self.thirdAxis = self._getAxis(2)
        
        self.invertedFirstAxis = '-' + self._getAxis(0)
        self.invertedSecondAxis = '-' + self._getAxis(1)
        self.invertedThirdAxis = '-' + self._getAxis(2)
        
        self.firstAxisList = self.axis[self.firstAxis]
        self.secondAxisList = self.axis[self.secondAxis]
        self.thirdAxisList = self.axis[self.thirdAxis]
        
        self.invertedFirstAxisList = self.axis['-' + self.firstAxis]
        self.invertedSecondAxisList = self.axis['-' + self.secondAxis]
        self.invertedThirdAxisList = self.axis['-' + self.thirdAxis]
        
    def _getAxis(self, axisIndex):
        return self.currentOrder[ axisIndex%3 ]
    
    def _getAxisArray(self, axis):
        return self.axis[ axis ]

class oyRigJoint(object):
    '''its a data class to hold the data about my joints\n
    it holds the data listed below:
    
    joint itself
    the baseLocator 
    the upLocator
    
    a rotateOrder object to hold :
        joint rotate order as str
        joint rotate order as int
        joint first axis as str
        joint second axis as str
        joint third axis as str
        joint first axis as int list
        joint second axis as int list
        joint third axis as int list
    '''
    
    
    
    def __init__(self, joint):
        self._joint = oyAux.getValidDagNode( joint )
        self._baseLocator = pm.Transform
        self._upLocator = pm.Transform
        
        self._rotateOrder = rotateOrder( self._joint.getAttr('ro') )
    
    
    
    def set_baseLocator(self, locator):
        self._baseLocator = oyAux.getValidDagNode( locator )
    
    
    
    def set_upLocator(self, locator):
        self._upLocator = oyAux.getValidDagNode( locator )
    
    
    
    def _saveSettings(self):
        """saves settings to the joints oyRigData attribute
        """
        self._createAttributes()
        
        # connect up and base locators
        pm.connectAttr( )
    
    
    
    def _createAttributes(self):
        """creates data attribute
        """
        
        #
        
        if not self._joint.hasAttr("oyRigData"):
            pm.addAttr(self._joint, ln="oyRigData", at="compound", nc=4)
        
        if not self._joint.hasAttr( "baseLocator" ):
            pm.addAttr( self._joint, ln="baseLocator", at="message", p="oyRigData" )
        
        if not self._joint.hasAttr( "upLocator" ):
            pm.addAttr( self._joint, ln="upLocator", at="message", p="oyRigData" )
    
    
    
    def
    
    
    